#include "MPU6500.hpp"
#include <chrono>
#include <thread>
#include <iostream>
#include <stdexcept>

#define write_wrapper(x) do { \
    if (!(x)) { \
        throw std::runtime_error("I2C write failed"); \
    } \
} while (0)


MPU6500::MPU6500(I2CDevice& i2cDevice) : i2c(i2cDevice) {
    accel_scale = 4.0f / 32768.0f;
    gyro_scale  = 500.0f / 32768.0f;
}

bool MPU6500::initialize() {
    write_wrapper(i2c.writeByte(0x6B, 0x00)); // Wake up
    std::this_thread::sleep_for(std::chrono::milliseconds(100));
    write_wrapper(i2c.writeByte(0x1B, 0x08)); // Gyro ±500°/s
    write_wrapper(i2c.writeByte(0x1C, 0x08)); // Accel ±4g
    write_wrapper(i2c.writeByte(0x1A, 0x03)); // DLPF ~44Hz
    write_wrapper(i2c.writeByte(0x19, 0x27)); // Sample divider = 39+1 = 40 -> 1kHz / 40 = 25Hz
    return true;
}

int16_t MPU6500::combineBytes(uint8_t high, uint8_t low) {
    union {
        int16_t value;
        uint8_t bytes[2];
    } combined;
    combined.bytes[0] = low;
    combined.bytes[1] = high;
    return combined.value;
}

void MPU6500::calibrate(int samples) {
    float ax_sum = 0, ay_sum = 0, az_sum = 0;
    float gx_sum = 0, gy_sum = 0, gz_sum = 0;

    float ax, ay, az, gx, gy, gz;
    
    for (int i = 0; i < samples; ++i) {
        readSensor(ax, ay, az, gx, gy, gz);

        ax_sum += ax;
        ay_sum += ay;
        az_sum += az - 1.0f;
        gx_sum += gx;
        gy_sum += gy;
        gz_sum += gz;

        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    }

    accel_offset[0] = ax_sum / samples;
    accel_offset[1] = ay_sum / samples;
    accel_offset[2] = az_sum / samples;
    gyro_offset[0] = gx_sum / samples;
    gyro_offset[1] = gy_sum / samples;
    gyro_offset[2] = gz_sum / samples;
}

void MPU6500::readSensor(float& ax, float& ay, float& az, float& gx, float& gy, float& gz) {
    uint8_t buffer[14];
    if (!i2c.readBytes(0x3B, buffer, 14)) {
        throw std::runtime_error("Failed to read sensor data");
    }

    ax = combineBytes(buffer[0], buffer[1]) * accel_scale - accel_offset[0];
    ay = combineBytes(buffer[2], buffer[3]) * accel_scale - accel_offset[1];
    az = combineBytes(buffer[4], buffer[5]) * accel_scale - accel_offset[2];
    gx = combineBytes(buffer[8], buffer[9]) * gyro_scale  - gyro_offset[0];
    gy = combineBytes(buffer[10], buffer[11]) * gyro_scale - gyro_offset[1];
    gz = combineBytes(buffer[12], buffer[13]) * gyro_scale - gyro_offset[2];
}
